from statistics import mode
import torch

from torch import nn, optim

from torch.autograd import Variable

from torch.utils.data import DataLoader

from torchvision import datasets, transforms

import cnn

import torchvision.models

from torchvision.models.resnet import resnet18, resnet34
from torchvision.models import vgg19, vgg19_bn, vgg16, inception_v3, mobilenet_v2
from PIL import Image
# 定义一些超参数

batch_size = 64

learning_rate = 0.2

num_epoches = 20



# 数据预处理。transforms.ToTensor()将图片转换成PyTorch中处理的对象Tensor,并且进行标准化（数据在0~1之间）

# transforms.Normalize()做归一化。它进行了减均值，再除以标准差。两个参数分别是均值和标准差

# transforms.Compose()函数则是将各种预处理的操作组合到了一起

data_tf = transforms.Compose(

    [
    #### 引入数据增强
    # torchvision.transforms.RandomAffine(degrees = 0,translate=(0.1, 0.1)),#平移距离
    # torchvision.transforms.RandomRotation((-10,10)),#将图片随机旋转（-10,10）度
    ### 完毕
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
    
    ])



# 数据集的下载器

train_dataset = datasets.MNIST(

    root='./data', train=True, transform=data_tf, download=True)

test_dataset = datasets.MNIST(root='./data', train=False, transform=data_tf)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)




# 选择模型

# model = cnn.NeuralNetwork()
# model = cnn.CNN()
model = cnn.VGG16(10)

if torch.cuda.is_available():

    model = model.cuda()



# 定义损失函数和优化器

criterion = nn.CrossEntropyLoss()

optimizer = optim.ASGD(model.parameters(), lr=learning_rate)



# 训练模型

epoch = 0

for data in train_loader:

    img, label = data

    if torch.cuda.is_available():

        img = img.cuda()

        label = label.cuda()

    else:

        img = Variable(img)

        label = Variable(label)

    out = model(img)

    loss = criterion(out, label)

    print_loss = loss.data.item()

    optimizer.zero_grad()

    loss.backward()

    optimizer.step()

    epoch+=1

    if epoch%50 == 0:

        print('item: {}, loss: {:.4}'.format(epoch, loss.data.item()))


# 保存和加载整个模型

torch.save(model, 'CNN_for_MNIST.pth')




# 模型评估

model.eval()

eval_loss = 0

eval_acc = 0

for data in test_loader:

    img, label = data
    img = Variable(img)

    if torch.cuda.is_available():

        img = img.cuda()
        label = label.cuda()

    out = model(img)
    loss = criterion(out, label)
    eval_loss += loss.data.item()*label.size(0)
    _, pred = torch.max(out, 1)
    num_correct = (pred == label).sum()
    eval_acc += num_correct.item()

print('Test Loss: {:.6f}, Acc: {:.6f}'.format(

    eval_loss / (len(test_dataset)),

    eval_acc / (len(test_dataset))

))
